package com.ruoyi.web.ai.listener;

import cn.hutool.json.JSONUtil;
import com.alibaba.fastjson.JSON;
import com.plexpt.chatgpt.entity.chat.ChatChoice;
import com.plexpt.chatgpt.entity.chat.ChatCompletionResponse;
import com.plexpt.chatgpt.entity.chat.Message;
import com.plexpt.chatgpt.listener.AbstractStreamListener;
import com.plexpt.chatgpt.util.SseHelper;
import com.ruoyi.common.ai.OpenAiResponse;
import com.ruoyi.common.base.ai.enums.CommonError;
import com.ruoyi.common.base.ai.vo.AIApiKeyCache;
import com.ruoyi.web.ai.AIDynamicKey;
import com.ruoyi.web.controller.ai.AIChatController;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.sse.EventSource;
import org.apache.logging.log4j.util.Strings;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.function.Consumer;

/**
 * 自定义SSE监听器
 *
 * @author gdk
 * @since 2023/5/17 14:59
 */
@Slf4j
public class CustomizeSseStreamListener extends AbstractStreamListener {
    /**
     * 账号被封了
     */
    private static final String ACCOUNT_DEACTIVATED = "account_deactivated";
    /**
     * key不正确
     */
    private static final String INVALID_API_KEY = "invalid_api_key";

    final SseEmitter sseEmitter;
    final AIDynamicKey aiDynamicKey;
    final AIApiKeyCache currentKey;
    protected Consumer<?> onOpen;

    public CustomizeSseStreamListener(SseEmitter sseEmitter, AIDynamicKey aiDynamicKey, AIApiKeyCache currentKey) {
        this.sseEmitter = sseEmitter;
        this.aiDynamicKey = aiDynamicKey;
        this.currentKey = currentKey;
    }

    /**
     * 事件处理
     *
     * @author gdk
     * @since 2023/5/21 16:09
     */
    @Override
    public void onEvent(EventSource eventSource, String id, String type, String data) {
        if (data.equals("[DONE]")) {
            this.onComplate.accept(this.lastMessage);
        } else {
            ChatCompletionResponse response = (ChatCompletionResponse) JSON.parseObject(data, ChatCompletionResponse.class);
            List<ChatChoice> choices = response.getChoices();
            if (choices != null && !choices.isEmpty()) {
                // 必须直接返回delta对象，返回delta.getContent()会导致前端无法获取到真正的响应字符
                Message delta = (choices.get(0)).getDelta();
                try {
                    sseEmitter.send(SseEmitter.event().id("message").data(delta));
                } catch (IOException e) {
                    eventSource.cancel();
                }
            }
        }
    }

    @Override
    public void onMsg(String message) {
    }

    @Override
    public void onFailure(EventSource eventSource, Throwable throwable, Response response) {
        // aiDynamicKey.disableKey(currentKey, "测试停用apiKey");

        // 用户端连接中断，发送sse消息时抛出异常
        if (!Objects.isNull(throwable)
                && ("stream was reset: CANCEL".equals(throwable.getMessage())
                || "ResponseBodyEmitter has already completed".equals(throwable.getMessage()))) {
            return;
        }

        try {
            if (Objects.isNull(response)) {
                AIChatController.sendError("未能连接到模型，请重试", sseEmitter);
                eventSource.cancel();
                return;
            }

            String errorMsg = response.body().string();
            if (Strings.isBlank(errorMsg)) {
                AIChatController.sendError("未能连接到模型，请重试", sseEmitter);
                eventSource.cancel();
                return;
            }

            // 模型超载
            if (errorMsg.contains("That model is currently overloaded with other requests.")) {
                AIChatController.sendError("模型负荷已超载，请稍后重试", sseEmitter);

                // 账号被封/处理请求时出错
            } else if (response.code() == CommonError.OPENAI_AUTHENTICATION_ERROR.getCode()
                    || response.code() == CommonError.OPENAI_SERVER_ERROR.getCode()) {

                AIChatController.sendError("未能连接到模型，请重试", sseEmitter);

                OpenAiResponse openAiResponse = JSONUtil.toBean(errorMsg, OpenAiResponse.class);
                String errorCode = openAiResponse.getError().getCode();
                // 账号被封或者key不正确就移除掉
                if (ACCOUNT_DEACTIVATED.equals(errorCode) || INVALID_API_KEY.equals(errorCode)) {
                    aiDynamicKey.disableKey(currentKey, errorMsg);
                }

                // 请求频繁
            } else if (response.code() == CommonError.OPENAI_LIMIT_ERROR.getCode()) {
                AIChatController.sendError("请求模型过于频繁，请稍后重试", sseEmitter);

                // 非官方定义的错误code
            } else {
                AIChatController.sendError("请求模型时发生异常，请重试", sseEmitter);
                OpenAiResponse openAiResponse = JSONUtil.toBean(errorMsg, OpenAiResponse.class);
                if (Objects.nonNull(openAiResponse.getError())) {
                    log.error("ChatGPT接口调用，未知异常：{}", openAiResponse.getError().getMessage());
                }
            }

        } catch (Throwable e) {
            log.error("ChatGPT接口失败处理异常，异常信息：{}", e.getMessage());
        } finally {
            eventSource.cancel();
        }
    }


    @Override
    public void onError(Throwable throwable, String response) {
        SseHelper.complete(this.sseEmitter);
    }

    public void setOnOpen(Consumer<?> onOpen) {
        this.onOpen = onOpen;
    }
}
